import sympy

# 注释由AI辅助完成
# 假定有S1, S2两个参考系，S2 相对 S1 向右以速度u运动
# 在S1 中观察到粒子“平常速度为” (vx1,vy1,vz1)，在S2 中观察到粒子“平常速度为”  (vx2,vy2,vz2)
# 假设已知S2中粒子的速度，求S1中粒子的速度
vx2, vy2, vz2 = sympy.symbols('vx2,vy2,vz2')
c = sympy.symbols('c', positive=True) # 光速
u = sympy.symbols('u', positive=True)

# 在参考系S2中，由粒子的“平常速度”构建4-速度 U2
# 4-速度的定义是 U = γ * [c, vx, vy, vz]
# 4-速度中，γ 由当前参考系观察到的粒子速度计算 γ = 1/sqrt(1 - (vx2²+vy2²+vz2²)/c²)
g2 = 1 / sympy.sqrt(1 - (vx2**2 + vy2**2 + vz2**2) / c**2)
U2 = g2 * sympy.Matrix([[c], [vx2], [vy2], [vz2]])

# 计算S2->S1参考系变换的Lorentz变换矩阵
g = 1 / sympy.sqrt(1 - u**2 / c**2)
b = u / c
L = sympy.Matrix([
    [g,    g*b,  0,    0],
    [g*b,  g,     0,    0],
    [0,     0,      1,    0],
    [0,     0,      0,    1]
])

# 应用Lorentz变换：将4-速度从S2系变换到S1系
U1 = L @ U2

# 在参考系S1中，由粒子的4-速度U1 读出其“平常速度”
# 由于4-速度是 U = γ * [c, vx, vy, vz]，因此先读出 γ
# 随后计算“平常速度”
g1 = U1[0] / c
vx1 = U1[1] / g1
vy1 = U1[2] / g1
vz1 = U1[3] / g1

# 打印简化后的速度分量表达式
# 这些就是相对论速度叠加公式的具体形式
print("S1系中的x方向速度分量:")
print(vx1.simplify())
print("\nS1系中的y方向速度分量:")
print(vy1.simplify())
print("\nS1系中的z方向速度分量:")
print(vz1.simplify())
